import torch
import torch.nn as nn
from ..config.config import Config

class RNN(nn.Module):
    def __init__(self):
        """
        构造一个自定义RNN模型，实现短文本分类
        :param input_size: 输入张量最后一个维度的大小
        :param hidden_size: 隐藏层张量最后一个维度的大小
        :param output_size: 输出层张量最后一个维度的大小
        """
        super(RNN, self).__init__()
        # 定义RNN
        self.rnn = nn.RNN(Config.INPUT_SIZE, Config.HIDDEN_SIZE, batch_first=True)
        # 线性层
        self.fc = nn.Linear(Config.HIDDEN_SIZE, Config.N_CATEGORIES)
        # 定义最后的softmax
        self.softmax = nn.LogSoftmax(dim=-1)

    def forward(self, input1, hidden1):
        """
        前向传播
        :param input1: x(t)
        :param hidden1: h(t-1)
        :return:
        """
        output, hidden = self.rnn(input1.unsqueeze(0), hidden1.unsqueeze(0))
        # 进入Linear
        output = self.fc(output.squeeze(0))
        # 进入softmax
        output = self.softmax(output)
        return output, hidden.unsqueeze(0)

    def init_hidden(self):
        # 初始化 [1, hidden_size] 全为0的张量
        return torch.zeros(1, Config.HIDDEN_SIZE)